package kodkod.analysis.ast;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.TreeSet;

import kodkod.analysis.Analyzer;
import kodkod.analysis.Problem;
import kodkod.ast.BinaryExpression;
import kodkod.ast.BinaryFormula;
import kodkod.ast.BinaryIntExpression;
import kodkod.ast.ComparisonFormula;
import kodkod.ast.Comprehension;
import kodkod.ast.ConstantExpression;
import kodkod.ast.ConstantFormula;
import kodkod.ast.Decl;
import kodkod.ast.Decls;
import kodkod.ast.ExprToIntCast;
import kodkod.ast.Expression;
import kodkod.ast.Formula;
import kodkod.ast.IfExpression;
import kodkod.ast.IfIntExpression;
import kodkod.ast.IntComparisonFormula;
import kodkod.ast.IntConstant;
import kodkod.ast.IntExpression;
import kodkod.ast.IntToExprCast;
import kodkod.ast.LeafExpression;
import kodkod.ast.MultiplicityFormula;
import kodkod.ast.NaryExpression;
import kodkod.ast.NaryFormula;
import kodkod.ast.NaryIntExpression;
import kodkod.ast.Node;
import kodkod.ast.NotFormula;
import kodkod.ast.ProjectExpression;
import kodkod.ast.QuantifiedFormula;
import kodkod.ast.Relation;
import kodkod.ast.RelationPredicate;
import kodkod.ast.SumExpression;
import kodkod.ast.UnaryExpression;
import kodkod.ast.UnaryIntExpression;
import kodkod.ast.Variable;
import kodkod.ast.operator.ExprCompOperator;
import kodkod.ast.operator.FormulaOperator;
import kodkod.ast.visitor.VoidVisitor;

public abstract class GraphVisitor extends ExtendedVoidVisitor {

	HashMap<Node, Integer> nodeIds = new HashMap<Node, Integer>();
	Integer idCounter = 0;
		
	public static boolean dropSelfEqualities = true;

	static public boolean negated = false;
		
	protected abstract void newCompound(Node parent, Node child);
	protected abstract void newEdge(Node parent, Node child, int pos);
	protected abstract void newEdge(Node parent, Node child);
	protected abstract void newNode(Node node);
	
	ArrayList<Node> currentPath = new ArrayList<Node>();
		
	protected void enter(Node node) {
		currentPath.add(node);
	}
	protected void leave(Node node) {
		currentPath.remove(node);
	}
		
	protected int getNodeId(Node node) {
		Integer id = nodeIds.get(node);
		if (id == null) {
			nodeIds.put(node, ++idCounter);
			newNode(node);
			return idCounter;
		} else {
			return id.intValue();
		}
	}
	
	public int numberOfNodes() {
		return idCounter;
	}
		
	@Override
	public void visit(Decls decls) {
		enter(decls);
		
		for (Decl decl : decls) {
			newEdge(decls, decl);			
			decl.accept(this);
		}
		
		leave(decls);
	}

	@Override
	public void visit(Decl decl) {
		enter(decl);
		
		Variable var = decl.variable();
		Expression exp = decl.expression();
		
//		InitialAnalysis.decls.put(decl.variable(), decl.expression());
		
		newEdge(decl, var);
		newEdge(decl, exp);
		var.accept(this);
		exp.accept(this);
		
		leave(decl);
	}

	@Override
	public void visit(Relation relation) {}

	@Override
	public void visit(Variable variable) {}

	@Override
	public void visit(ConstantExpression constExpr) {}

	@Override
	public void visit(UnaryExpression unaryExpr) {
		enter(unaryExpr);
		newEdge(unaryExpr, unaryExpr.expression());	
		unaryExpr.expression().accept(this);
		leave(unaryExpr);
	}

	@Override
	public void visit(BinaryExpression binExpr) {
		enter(binExpr);
		newEdge(binExpr, binExpr.left(), 1);
		newEdge(binExpr, binExpr.right(), 2);
		binExpr.left().accept(this);
		binExpr.right().accept(this);
		leave(binExpr);
	}

	@Override
	public void visit(NaryExpression expr) {
		enter(expr);
		for(Expression child : expr) {
			newEdge(expr, child);
			child.accept(this);
		}
		leave(expr);
	}

	@Override
	public void visit(Comprehension comprehension) {
		enter(comprehension);
		newEdge(comprehension, comprehension.formula());
		if (comprehension.decls() instanceof Decl) {
			Decl decl = (Decl)comprehension.decls();
			newEdge(comprehension, decl.variable());
			decl.variable().accept(this);
			newEdge(decl.variable(), decl.expression());
			decl.expression().accept(this);
		} else if (comprehension.decls() instanceof Decls) {
			for (Decl decl : comprehension.decls()) {
				newEdge(comprehension, decl.variable());
				decl.variable().accept(this);
				newEdge(decl.variable(), decl.expression());
				decl.expression().accept(this);
			}
		}
		comprehension.formula().accept(this);
		leave(comprehension);
	}

	@Override
	public void visit(IfExpression ifExpr) {
		enter(ifExpr);
		newEdge(ifExpr, ifExpr.condition());
		newEdge(ifExpr, ifExpr.thenExpr(), 3);
		newEdge(ifExpr, ifExpr.elseExpr(), 4);
		ifExpr.condition().accept(this);
		ifExpr.thenExpr().accept(this);
		ifExpr.elseExpr().accept(this);
		leave(ifExpr);
	}

	@Override
	public void visit(ProjectExpression project) {
		enter(project);
		newEdge(project, project.expression());
		project.expression().accept(this);
		for(int i = 0, arity = project.arity(); i < arity; i++) {
			newEdge(project, project.column(i));
			project.column(i).accept(this);
		}
		leave(project);
	}

	@Override
	public void visit(IntToExprCast castExpr) {
		enter(castExpr);
		newEdge(castExpr, castExpr.intExpr());
		castExpr.intExpr().accept(this);
		leave(castExpr);
	}

	@Override
	public void visit(IntConstant intConst) {}

	@Override
	public void visit(ExprToIntCast intExpr) {
		enter(intExpr);
		newEdge(intExpr, intExpr.expression());
		intExpr.expression().accept(this);
		leave(intExpr);
	}

	@Override
	public void visit(IfIntExpression intExpr) {
		enter(intExpr);
		newEdge(intExpr, intExpr.condition());
		newEdge(intExpr, intExpr.thenExpr(), 1);
		newEdge(intExpr, intExpr.elseExpr(), 2);
		intExpr.condition().accept(this);
		intExpr.thenExpr().accept(this);
		intExpr.elseExpr().accept(this);
		leave(intExpr);
	}

	@Override
	public void visit(NaryIntExpression intExpr) {
		enter(intExpr);
		for(IntExpression child : intExpr) {
			newEdge(intExpr, child);
			child.accept(this);
		}
		leave(intExpr);
	}

	@Override
	public void visit(BinaryIntExpression intExpr) {
		enter(intExpr);
		newEdge(intExpr, intExpr.left(), 1);
		newEdge(intExpr, intExpr.right(), 2);
		intExpr.left().accept(this);
		intExpr.right().accept(this);
		leave(intExpr);
	}

	@Override
	public void visit(UnaryIntExpression intExpr) {
		enter(intExpr);
		newEdge(intExpr, intExpr.intExpr());
		intExpr.intExpr().accept(this);
		leave(intExpr);
	}

	@Override
	public void visit(SumExpression intExpr) {
		enter(intExpr);
		newEdge(intExpr, intExpr.decls());
		newEdge(intExpr, intExpr.intExpr());
		intExpr.decls().accept(this);
		intExpr.intExpr().accept(this);
		leave(intExpr);
	}

	@Override
	public void visit(IntComparisonFormula intComp) {
		enter(intComp);
		newEdge(intComp, intComp.left(), 3);
		newEdge(intComp, intComp.right(), 4);
		intComp.left().accept(this);
		intComp.right().accept(this);
		leave(intComp);
	}

	@Override
	public void visit(QuantifiedFormula qFormula) {
		enter(qFormula);
		newEdge(qFormula, qFormula.formula());
		qFormula.formula().accept(this);
		
		if (qFormula.decls() instanceof Decl) {
			Decl decl = (Decl)qFormula.decls();
			newEdge(qFormula, decl.variable());
			decl.variable().accept(this);
			newEdge(decl.variable(), decl.expression());
			//newCompound(decl.expression(), decl.variable());
			decl.expression().accept(this);
		} else if (qFormula.decls() instanceof Decls) {
			for (Decl decl : qFormula.decls()) {
				newEdge(qFormula, decl.variable());
				decl.variable().accept(this);
				newEdge(decl.variable(), decl.expression());
				//newCompound(decl.expression(), decl.variable());
				decl.expression().accept(this);
			}
		}
		leave(qFormula);
	}

	@Override
	public void visit(NaryFormula node) {
		enter(node);
		if (node.op().equals(FormulaOperator.AND)) {
			ArrayList<Formula> joinedFormulas = collectAllFormulasByType(node, FormulaOperator.AND);
			for (Formula f : joinedFormulas) {
				if (dropSelfEqualities && f instanceof ComparisonFormula) {
					ComparisonFormula cf = (ComparisonFormula)f;
					if (cf.left() == cf.right()) continue;
				}
				
				if (f instanceof RelationPredicate.TotalOrdering) {
					f.accept(this);
					continue; // no extra node for Ordering
				}

				f.accept(this);
				newEdge(node, f);
			}
		}
		
		else if (node.op().equals(FormulaOperator.OR)) {			
			ArrayList<Formula> joinedFormulas = collectAllFormulasByType(node, FormulaOperator.OR);
			for (Formula f : joinedFormulas) {
				if (dropSelfEqualities && f instanceof ComparisonFormula) {
					ComparisonFormula cf = (ComparisonFormula)f;
					if (cf.left() == cf.right()) continue;
				}
				
				if (f instanceof RelationPredicate.TotalOrdering) {
					f.accept(this);
					continue; // no extra node for Ordering
				}

				f.accept(this);
				newEdge(node, f);
			}
		}
		leave(node);
	}

	@Override
	public void visit(BinaryFormula node) {
		enter(node);
		if (node.op().equals(FormulaOperator.AND)) {
			ArrayList<Formula> joinedFormulas = collectAllFormulasByType(node, FormulaOperator.AND);
			for (Formula f : joinedFormulas) {
				newEdge(node, f);
				f.accept(this);
			}
		}
		
		else if (node.op().equals(FormulaOperator.OR)) {
			ArrayList<Formula> joinedFormulas = collectAllFormulasByType(node, FormulaOperator.OR);
			for (Formula f : joinedFormulas) {
				newEdge(node, f);
				f.accept(this);
			}
		}
		
		else if (node.op().equals(FormulaOperator.IMPLIES)) {
			newEdge(node, node.left(), 3);
			newEdge(node, node.right(), 4);
			
			boolean former = negated;
			negated = !negated;
			node.left().accept(this);
			negated = former;
			
			node.right().accept(this);
		}
		
		else if (node.op().equals(FormulaOperator.IFF)) {
			newEdge(node, node.left(), 3);
			newEdge(node, node.right(), 4);
			node.left().accept(this);
			node.right().accept(this);
		}
		leave(node);
	}
	
	private ArrayList<Formula> collectAllFormulasByType(BinaryFormula node, FormulaOperator optype) {
		ArrayList<Formula> list = new ArrayList<Formula>();
		
		Formula check = node.left();
				
		if (check instanceof BinaryFormula && ((BinaryFormula)check).op().equals(optype)) {
			list.addAll(collectAllFormulasByType((BinaryFormula)check, optype));
		} else {
			list.add(check);
		}
		
		check = node.right();
				
		if (check instanceof BinaryFormula && ((BinaryFormula)check).op().equals(optype)) {
			list.addAll(collectAllFormulasByType((BinaryFormula)check, optype));
		} else {
			list.add(check);
		}
		
		return list;
	}
	
	private ArrayList<Formula> collectAllFormulasByType(NaryFormula node, FormulaOperator optype) {
		ArrayList<Formula> list = new ArrayList<Formula>();
		
		for (Formula child : node) {
			if (child instanceof BinaryFormula && ((BinaryFormula)child).op().equals(optype)) {
				list.addAll(collectAllFormulasByType((BinaryFormula)child, optype));
			} else {
				list.add(child);
			}
		}
		
		return list;
	}

	@Override
	public void visit(NotFormula not) {		
		enter(not);
		newEdge(not, not.formula());

		boolean former = negated;
		negated = !negated;
		not.formula().accept(this);
		negated = former;
		
		leave(not);
	}

	@Override
	public void visit(ConstantFormula constant) {}

	@Override
	public void visit(ComparisonFormula compFormula) {
		enter(compFormula);
		if (compFormula.op() != ExprCompOperator.EQUALS) {
			newEdge(compFormula, compFormula.left(), 3);
			newEdge(compFormula, compFormula.right(), 4);
		} else {
			newEdge(compFormula, compFormula.left());
			newEdge(compFormula, compFormula.right());
		}
		compFormula.left().accept(this);
		compFormula.right().accept(this);
		leave(compFormula);
	}

	@Override
	public void visit(MultiplicityFormula multFormula) {
		enter(multFormula);
		newEdge(multFormula, multFormula.expression());
		multFormula.expression().accept(this);
		leave(multFormula);
	}

	@Override
	public void visit(RelationPredicate predicate) {
		enter(predicate);
		predicate.relation().accept(this);
		if (predicate.name()==RelationPredicate.Name.FUNCTION) {
			newEdge(predicate, predicate.relation());
			final RelationPredicate.Function fp = (RelationPredicate.Function) predicate;
			newEdge(predicate, fp.domain());
			newEdge(predicate, fp.range());
			fp.domain().accept(this);
			fp.range().accept(this);
		} else if (predicate.name()==RelationPredicate.Name.TOTAL_ORDERING) {
			final RelationPredicate.TotalOrdering tp = (RelationPredicate.TotalOrdering) predicate;
			newCompound(tp.ordered(), tp.relation());
			newCompound(tp.ordered(), tp.first());
			newCompound(tp.ordered(), tp.last());
		}
		leave(predicate);
	}
}
